from typing import List, Dict, Optional
import numpy as np


class Block:
    """内存块"""

    def __init__(self, block_id: int, block_size: int):
        self.block_id = block_id
        self.block_size = block_size
        self.ref_count = 0
        self.token_ids: List[int] = []

    def can_append(self, num_tokens: int) -> bool:
        """检查是否可以添加token"""
        return len(self.token_ids) + num_tokens <= self.block_size

    def append(self, token_ids: List[int]):
        """添加token"""
        self.token_ids.extend(token_ids)

    def is_full(self) -> bool:
        """检查块是否已满"""
        return len(self.token_ids) >= self.block_size


class BlockAllocator:
    """块分配器"""

    def __init__(self, num_blocks: int, block_size: int):
        self.num_blocks = num_blocks
        self.block_size = block_size
        self.free_blocks: List[Block] = []
        self.allocated_blocks: Dict[int, Block] = {}

        # 初始化空闲块
        for i in range(num_blocks):
            self.free_blocks.append(Block(i, block_size))

    def allocate(self, num_tokens: int) -> List[Block]:
        """分配块"""
        num_blocks_needed = (num_tokens + self.block_size - 1) // self.block_size

        if len(self.free_blocks) < num_blocks_needed:
            raise RuntimeError("Not enough free blocks")

        allocated = []
        for _ in range(num_blocks_needed):
            block = self.free_blocks.pop()
            block.ref_count = 1
            allocated.append(block)
            self.allocated_blocks[block.block_id] = block

        return allocated

    def free(self, blocks: List[Block]):
        """释放块"""
        for block in blocks:
            block.ref_count -= 1
            if block.ref_count == 0:
                block.token_ids = []
                self.free_blocks.append(block)